from keras.preprocessing import image as kim
import cv2
def readImage(filename, shape, flag):
    if shape==None:
        m = cv2.imread(filename, flag)
    else:
        m = cv2.resize( cv2.imread(filename, flag) ,shape)
    if len(m.shape) < 3:
        m.resize(m.shape[0],m.shape[1],1)
    return m

class MyImageDataGenerator(kim.ImageDataGenerator):
    def __init__(self, featurewise_center = False, samplewise_center = False, featurewise_std_normalization = False, samplewise_std_normalization = False, zca_whitening = False, zca_epsilon = 1e-06, rotation_range = 0.0, width_shift_range = 0.0, height_shift_range = 0.0, shear_range = 0.0, zoom_range = 0.0, channel_shift_range = 0.0, fill_mode = 'nearest', cval = 0.0, horizontal_flip = False, vertical_flip = False, rescale = None, preprocessing_function = None, data_format = None):
        kim.ImageDataGenerator.__init__(self, featurewise_center, samplewise_center, featurewise_std_normalization, samplewise_std_normalization, zca_whitening, zca_epsilon, rotation_range, width_shift_range, height_shift_range, shear_range, zoom_range, channel_shift_range, fill_mode, cval, horizontal_flip, vertical_flip, rescale, preprocessing_function, data_format)

    def flow_from_list(self, trainX, trainY, batch_size = 32, reshape = None, flags = None,seed = None):
        LEN = len(trainY)
        kim.np.random.seed(seed)
        k = 0
        while 1:
            X = kim.np.array( [ readImage(f,reshape,flags) for f in trainX[k*batch_size : (k+1)*batch_size] ] )
            Y = trainY[k*batch_size : (k+1)*batch_size,:]

            f = self.flow(X, Y, batch_size, shuffle = False,seed = seed)
            
            yield f.next()
            k+=1
            if (k*batch_size>=LEN):
                k = 0


